import argparse
import os

import torch
from thop import profile

from configs.cmnext_init_cfg import _C as config, update_config
from models.cmnext_conf_3 import CMNeXtWithConf

available_device = 1
parser = argparse.ArgumentParser(description='')
parser.add_argument('-gpu', '--gpu', type=int, default=available_device, help='device, use -1 for cpu')
parser.add_argument('-log', '--log', type=str, default='INFO', help='logging level')
parser.add_argument('-train_bayar', '--train_bayar', default=False, action='store_true', help='finetune bayar conv')
parser.add_argument('-exp', '--exp', type=str, default='experiments/ec_example.yaml', help='Yaml experiment file')
parser.add_argument('opts', help="other options", default=None, nargs=argparse.REMAINDER)

os.environ['CUDA_VISIBLE_DEVICES'] = '{}'.format(available_device)

args = parser.parse_args()

config = update_config(config, args.exp)
model = CMNeXtWithConf(config.MODEL)


input1_1 = torch.randn(1, 3, 512, 512)
input1_2 = torch.randn(1, 3, 512, 512)
input1_3 = torch.randn(1, 3, 512, 512)
input1_4 = torch.randn(1, 3, 512, 512)
input1 = [input1_1,input1_2,input1_3,input1_4]
input2 = torch.randn(1, 1, 512, 512)

flops, params = profile(model, inputs=(input1, 1,input2))
print('flops:{}'.format(flops))

print('params:{}'.format(params))
